from prestool import Tool
from kazoo.client import KazooClient
import random
import re


class ZooKeeperUtil(Tool.Tool):
    def __init__(self):
        super().__init__()
        self.zk = KazooClient(hosts='zk.xxxx.cn:2181')  # 需要替换zk的地址
        self.zk.start()

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.zk.stop()
        self.zk.close()

    # 获取API列表
    def get_service_list(self):
        return self.zk.get_children("/dubbo")

    # 获取dubbo接口的信息
    def get_dubbo_info(self, _service):
        providers = self.zk.get_children(f'/dubbo/{_service}/providers')
        if providers:
            path, temp = self.url_decode(random.choice(providers)).split('dubbo://')[1].split('/')
            host, port = path.split(':')
            return {
                'host': host,
                'port': port,
                'service': temp.split('?')[0],
                'method': temp.split('methods=')[1].split('&')[0].split(',')
            }


class ToolUtil(Tool.Tool):
    def __init__(self):
        super().__init__()

    # 请求dubbo接口
    def dubbo_request(self, host, port, service, method, args):
        res = self.invoke_dubbo(host, port, service, method, args)[0]
        t = (re.findall("elapsed: (.*).\r", res) or [""])[0]
        result = self.json_loads(res.replace(f'\nelapsed: {t}.', '').
                                 replace(f'Use default service {service}.\r\nresult: ', '').
                                 replace('\n', '').replace('\r', ''))
        return self.json_dumps(result, indent=2)

    # 请求dubbo接口
    def dubbo_client(self, service, method, args):
        with ZooKeeperUtil() as zk:
            info = zk.get_dubbo_info(service)
            host, port = info['host'], info['port']
            return self.dubbo_request(host, port, service, method, args), info


# 获取dubbo接口 API 列表
def get_service_list(zk_ele):
    return zk_ele.get_service_list()


# 异常兼容
def get_info_error(zk_ele, service):
    info = zk_ele.get_dubbo_info(service) if service else None
    error = None if info else f'{service}\n服务异常,请重试'
    return info, error